#!/usr/bin/env python

import sys
import glob
import os
import xseg.automaticSeg as automaticSeg
import xseg.adaSeg adaSeg
import xseg.bone_extract as bone_extract
import xseg.vas_seg as vas_seg
import itk
import numpy as np
import scipy.ndimage as ndi
import lung_convex_hull

def getparser():
    parser = optionparser()
    parser.add_option("-i","--img",dest='fname',
                      help='name of image to read')
    parser.add_option("-t","--thresh",dest='threshold',nargs=2,default=(0,10000),type='int')
    parser.add_option("-s","--seed",dest='seed',nargs=3,default=(0,0,0),
                      type='int')
    parser.add_option("-v", "--variance", dest='var', default=1.0, type = 'float')
    parser.add_option("-m", "--mode", dest = 'mode', nargs = 5, default = 1, type = 'int')

    return parser

def count(filename):
    reader = itk.ImageFileReader.ISS3.New()
    reader.SetFileName(filename)
    reader.Update()
    img = reader.GetOutput()
    array = itk.PyBuffer.ISS3.GetArrayFromImage(img)
    count = array.shape[0] * array.shape[1]* array.shape[2]
    return count

def adathresh(fname, num):
    low  = 0
    high = 400
    value = 100 
    tmp = os.path.splitext(fname)
    scale = count(fname)*1.0/92536832
    print scale
    while True : 
        outputfile = "rough_seg.nii.gz"
        size = vas_seg.vas_seg(fname, value, 1000, 1, num, outputfile)
        print size[0:5], value
        if size[0]< 2.8*10**6*scale : 
            high = value
            value = (high + low )/2
        elif size[0] > 3.2*10**6 *scale: 
            low = value
            value = (high + low)/2
        else: 
            break
        if (high-low) <5 :
            break
    return value

def extract_bone(fname, thresh, num):
    dis = (400 - thresh)*1.0/num 
    val = []
    stat = []
    tmp = os.path.splitext(fname)
    for i in range(0, num):
         thresh1 = thresh + int(dis*i)   
         outputfile =  tmp[0] + "__" + str(thresh1) + tmp[1]
         size1 = vas_seg.vas_seg(fname, thresh1, 1000, 2, 2, outputfile)
         stat1 = vas_seg.center(outputfile) 
         val.append( size1[1]*1.0*max(stat1)/min(stat1))
         stat.append( stat1)
         
    print val
    print stat 
    outputfile = tmp[0] + "__" + str(thresh + int(dis*val.index(max(val)))) + tmp[1]
    return outputfile 

def exclude_bone(fname, fbone):
    	reader = itk.ImageFileReader.ISS3.New()
        reader.SetFileName(fname)
        img1 = reader.GetOutput()
        arr1 = itk.PyBuffer.ISS3.GetArrayFromImage(img1)
        reader2 = itk.ImageFileReader.IUC3.New()
        reader2.SetFileName(fbone)
        img2 = reader2.GetOutput()
        arr2 = itk.PyBuffer.IUC3.GetArrayFromImage(img2) 
        idx = np.where(arr2>0)
        print len(idx)
        arr1[idx] = 0 
        img = itk.PyBuffer.ISS3.GetImageFromArray(arr1)
        outputfile = "bone_excluded.nii.gz"
        writer = itk.ImageFileWriter.ISS3.New()
        writer.SetInput(img)
        writer.SetFileName(outputfile)  
        writer.Update()
        return outputfile

def convex_hull(fbone): 
        reader = itk.ImageFileReader.IUC3.New()
        reader.SetFileName(fbone)
        reader.Update()
        img = itk.PyBuffer.IUC3.GetArrayFromImage(reader.GetOutput()) 
        newImg = np.zeros(img.shape,np.uint8)
        for i in xrange(newImg.shape[0]):
            ns = ndi.binary_fill_holes(img[i,:,:])
            newImg[i,:,:] = ns
    
        for i in xrange(newImg.shape[0]):
           for j in xrange(newImg.shape[1]):
              idx = np.where(img[i, j, :] > 0 )
              if len(idx[0]) > 0: 
                  index = range(min(idx[0]),max(idx[0]))
                  newImg[i, j, index] = 1  
           for j in xrange(newImg.shape[2]):
              idx = np.where(img[i, :, j] > 0 )
              if len(idx[0]) > 0: 
                 index = range(min(idx[0]),max(idx[0]))
                 newImg[i, index, j] = 1  
           for j in xrange(newImg.shape[1]):
              idx = np.where(img[i, j, :] > 0 )
              if len(idx[0]) > 0: 
                 index = range(min(idx[0]),max(idx[0]))
                 newImg[i, j, index] = 1
        newImg_out = itk.PyBuffer.IUC3.GetImageFromArray(newImg)
        writer = itk.ImageFileWriter.IUC3.New()
        writer.SetFileName(fbone)
        writer.SetInput(newImg_out)
        writer.Update()

def addImage(fname1, fname2, fout) :
     reader1 = itk.ImageFileReader.ISS3.New()
     reader1.SetFileName(fname1)
     reader2 = itk.ImageFileReader.ISS3.New()
     reader2.SetFileName(fname2)
     adder = itk.AddImageFilter.ISS3ISS3ISS3.New()
     adder.SetInput1(reader1.GetOutput())
     adder.SetInput2(reader2.GetOutput())
     writer = itk.ImageFileWriter.ISS3.New()
     writer.SetFileName(fout)
     writer.SetInput(adder.GetOutput())
     writer.Update()              
def run():
   presentDir = "/HD2/xiaofei/Code2"
   os.chdir(presentDir)
   dirs = glob.glob("PE00000") 
   print dirs
   fo = open("seg_result.txt", "w")

   for d in dirs:
      os.chdir(str(d))
      os.system("pwd")
      fname = str(d) + ".nii.gz" 
      print fname
      lung = lung_convex_hull.lungExtraction(fname)
      convexHull = lung_convex_hull.convexHullExtraction(lung)
      aftermask = lung_convex_hull.applyMask(fname, convexHull)
      aftermask = fname
      ret_val = adathresh(aftermask, 4)
      print "done"
      outputfile = extract_bone(aftermask, ret_val, 5)
      print outputfile, ret_val
      fbone = lung_convex_hull.convexHullExtraction(outputfile)
      output = exclude_bone(aftermask, fbone) 
      file2 = str(d) + "_seg1.nii.gz" 
      vas_seg.vas_seg(output, ret_val, 1200,1, 1, file2 )
      print outputfile, fname, aftermask 
      file1 = automaticSeg.seg(fname)
      print file1
      file3 = str(d) + "_seg.nii.gz"
      addImage(file1, file2, file3)
      cmd = "iterVote " + file3 + " " + str(d) + "_hole_filling.nii.gz" + " " + "1" + "  " + "10"
      print cmd
      os.system(cmd)
      fo.write( outputfile)
      fo.write("\n")
      fo.write(str( ret_val))
      os.chdir(presentDir)
  
def main():
   run()
if __name__ == '__main__':
    main()
